#include "stdafx.h"
#include "UnkHook.h"
#include "FixedMemMgr.h"

HRESULT CUnkHook::CreateInstance(IUnknown* pUnk, IQIARHook* pQIHook, UnkHookFlags uhFlags, bool fUseARHook, UnknownHook** ppHook)
{
    HRESULT hr;
    if (SUCCEEDED(hr = (*ppHook = new CUnkHook()) ? NOERROR : E_OUTOFMEMORY))
    {
        CUnkHook& Hook = **((CUnkHook**)ppHook);
        Hook.AddRef();
        Hook.m_Hook.HookUnk(reinterpret_cast<IProvideClassInfo*>(pUnk), pQIHook, uhFlags, false, fUseARHook, &Hook);
    }
    return hr;
}
void CBaseUnkHook::CHook::HookUnk(IProvideClassInfo* pUnk, IQIARHook* pQIHook, UnkHookFlags uhFlags, bool fWeakRefQIHook, bool fUseARHook, CBaseUnkHook* pOwner)
{
    m_fAllFlags = false;
    m_VTable[0] = reinterpret_cast<VTBL_ENTRY>(QueryInterface);
    m_VTable[3] = reinterpret_cast<VTBL_ENTRY>(GetClassInfo);
    if (fUseARHook)
    {
        m_VTable[1] = reinterpret_cast<VTBL_ENTRY>(AddRefHooked);
        m_VTable[2] = reinterpret_cast<VTBL_ENTRY>(ReleaseHooked);
        m_fARHook = true;
        m_fReportAddRef = (uhFlags & uhAddRef) == uhAddRef;
        m_fReportRelease = (uhFlags & uhRelease) == uhRelease;
    }
    else
    {
        m_VTable[1] = reinterpret_cast<VTBL_ENTRY>(AddRef);
        m_VTable[2] = reinterpret_cast<VTBL_ENTRY>(Release);
    }
    m_pOuter = pUnk; // Not AddRefed
    m_lpVTableStart = *((void**)pUnk);
    m_pQIHook = NULL;
    // Proceed even if we don't have flags now, they can be
    // set later.
    m_pQIHook = pQIHook;
    m_fHookQIBefore = (uhFlags & uhBeforeQI) == uhBeforeQI;
    m_fHookQIAfter = (uhFlags & uhAfterQI) == uhAfterQI;
    m_fCallMapIID = (uhFlags & uhMapIIDs) == uhMapIIDs;
    if (!(m_fWeakQIHookExplicit = fWeakRefQIHook))
    {
        // fWeakRefQIHook is true if the object deriving from
        // CBaseUnkHook also implements the passed in QIHook.
        // If this isn't the case, then we need to make sure
        // that the QIHook is not on the same object as the
        // pass in controlling IUnknown.
        IUnknown* pTest = NULL;
        pQIHook->QueryInterface(IID_IUnknown, (void**)&pTest);
        pQIHook->AddRef();
        if (m_fWeakQIHook = pTest == m_pOuter)
        {
            // Release the ref on the controlling IUnknown,
            // but keep the one on the QIHook itself.
            pTest->Release();
        }
        if (pTest)
        {
            pTest->Release();
        }
    }
    m_pOwner = pOwner;
    *((void**)m_pOuter) = (void*)this; 
}
void CBaseUnkHook::CHook::Unhook()
{
    if (m_pOuter)
    {
        *((void**)m_pOuter) = m_lpVTableStart;
        if (m_pQIHook && !m_fWeakQIHookExplicit)
        {
            if (m_fWeakQIHook)
            {
                m_pOuter->AddRef();
            }
            m_pQIHook->Release();
        }
        m_pOuter = NULL;
    }
}
STDMETHODIMP
CBaseUnkHook::get_Flags(UnkHookFlags *puhFlags)
{
    *puhFlags = (UnkHookFlags)((int)m_Hook.m_fAllFlags & (uhBeforeQI | uhAfterQI | uhMapIIDs | uhAddRef | uhRelease));
    return NOERROR;
}
STDMETHODIMP
CBaseUnkHook::put_Flags(UnkHookFlags uhFlags)
{
    m_Hook.m_fHookQIBefore = (uhFlags & uhBeforeQI) == uhBeforeQI;
    m_Hook.m_fHookQIAfter = (uhFlags & uhAfterQI) == uhAfterQI;
    m_Hook.m_fCallMapIID = (uhFlags & uhMapIIDs) == uhMapIIDs;
    if (m_Hook.m_fARHook)
    {
        m_Hook.m_fReportAddRef = (uhFlags & uhAddRef) == uhAddRef;
        m_Hook.m_fReportRelease = (uhFlags & uhRelease) == uhRelease;
    }
    return NOERROR;
}
HRESULT _stdcall CBaseUnkHook::CHook::QueryInterface(CHook** ppThis, REFIID riid, void** ppvObj)
{
    HRESULT hr;
    CHook& Hook = **ppThis;
    IID* piidUse = const_cast<IID*>(&riid);
    IID IIDMapped;
    *((void**)Hook.m_pOuter) = Hook.m_lpVTableStart;

    if ((int)Hook.m_fAllFlags & (uhBeforeQI | uhAfterQI | uhMapIIDs))
    {
        // We have to protect our object here because we're making a callback
        // on the object which owns us, which means that the hook variable could
        // be released, which releases us.  Make sure that we're still alive
        // after the call by addrefing.  Use the inline _Base functions instead
        // of the virtual functions for speed.
        Hook.m_pOwner->_BaseAddRef();
        IUnknown* pUnk = NULL;
        if (Hook.m_fCallMapIID)
        {
            IIDMapped = riid;
            if (FAILED(Hook.m_pQIHook->MapIID((VBGUID*)&IIDMapped)))
            {
                goto MapIIDCleanup;
            }
            piidUse = &IIDMapped;
        }
        if (Hook.m_fHookQIBefore)
        {
            // Try to get a result object.  If we get something out of the before
            // hook, then we don't do the normal QI.
            Hook.m_pQIHook->QIHook((VBGUID*)piidUse, uhBeforeQI, &pUnk, (IUnknown*)Hook.m_pOuter);
        }
        if (pUnk == NULL)
        {
            Hook.m_pOuter->QueryInterface(*piidUse, (void**)&pUnk);
        }
        if (Hook.m_fHookQIAfter)
        {
            // If we've succeeded, then pUnk is already set.  To block the QI,
            // the after call should simply free pUnk.
            Hook.m_pQIHook->QIHook((VBGUID*)piidUse, uhAfterQI, &pUnk, (IUnknown*)Hook.m_pOuter);
        }
MapIIDCleanup:
        // If we have an object, then we succeeded.  Otherwise, we don't
        // support the interface.
        hr = (*ppvObj = (void*)pUnk) ? NOERROR : E_NOINTERFACE;
        if (0 == Hook.m_pOwner->_BaseRelease())
        {
            // The hook is dead, just get out of here
            return hr;
        }
    }
    else
    {
        hr = Hook.m_pOuter->QueryInterface(*piidUse, ppvObj);
    }

    // We now need a special check to see if an interface
    // was returned on the main interface for an IID other
    // than the two we support in the override.  This shouldn't
    // happen outside the debugger, but we need to check because
    // it's an instant crash and burn.
    if (SUCCEEDED(hr) &&
        (*ppvObj == (void*)Hook.m_pOuter))
    {
        bool fBlock;
        switch (piidUse->Data1)
        {
            case 0: //IID_IUnknown.Data1:
                fBlock = !IsEqualIIDIgnoreData1(*piidUse, IID_IUnknown);
                break;
            case 0xB196B283: //IID_IProvideClassInfo.Data1:
                fBlock = !IsEqualIIDIgnoreData1(*piidUse, IID_IProvideClassInfo);
                break;
            default:
                fBlock = true;
                break;
        }
        if (fBlock)
        {
            Hook.m_pOuter->Release();
            *ppvObj = NULL;
            hr = E_NOINTERFACE;
        }
    }
    *((void**)Hook.m_pOuter) = (void*)&Hook; 
    return hr;
}
ULONG _stdcall CBaseUnkHook::CHook::AddRef(CHook** ppThis)
{
    ULONG Refs;
    CHook& Hook = **ppThis;
    *((void**)Hook.m_pOuter) = Hook.m_lpVTableStart; 
    Refs = Hook.m_pOuter->AddRef();
    *((void**)Hook.m_pOuter) = (void*)&Hook; 
    return Refs;
}
ULONG _stdcall CBaseUnkHook::CHook::AddRefHooked(CHook** ppThis)
{
    ULONG Refs;
    CHook& Hook = **ppThis;
    *((void**)Hook.m_pOuter) = Hook.m_lpVTableStart; 
    Refs = Hook.m_pOuter->AddRef();
    if (Hook.m_fReportAddRef)
    {
        // See notes on AddRef in CBaseUnkHookCHook::QueryInterface
        Hook.m_pOwner->_BaseAddRef();
        Hook.m_pQIHook->AfterAddRef((long)Refs, (long)ppThis);
        if (0 == Hook.m_pOwner->_BaseRelease())
        {
            // The hook is dead, just get out of here
            return Refs;
        }
    }
    *((void**)Hook.m_pOuter) = (void*)&Hook; 
    return Refs;
}
ULONG _stdcall CBaseUnkHook::CHook::Release(CHook** ppThis)
{
    ULONG Refs;
    CHook& Hook = **ppThis;
    *((void**)Hook.m_pOuter) = Hook.m_lpVTableStart; 
    if (Refs = Hook.m_pOuter->Release())
    {
        // Only replace if object isn't dead
        *((void**)Hook.m_pOuter) = (void*)&Hook;
    }
    return Refs;
}
ULONG _stdcall CBaseUnkHook::CHook::ReleaseHooked(CHook** ppThis)
{
    ULONG Refs;
    CHook& Hook = **ppThis;
    *((void**)Hook.m_pOuter) = Hook.m_lpVTableStart;
    if (Hook.m_fReportRelease)
    {
        if (Hook.m_fWeakQIHook || Hook.m_fWeakQIHookExplicit)
        {
            // m_pQIHook is on the same object as the hooked unknown.
            // If the object dies, then there is no way to report
            // the last release.
            if (Refs = Hook.m_pOuter->Release())
            {
                // See notes on AddRef in CBaseUnkHookCHook::QueryInterface
                Hook.m_pOwner->_BaseAddRef();
                Hook.m_pQIHook->AfterRelease((long)Refs, (long)ppThis);
                if (0 == Hook.m_pOwner->_BaseRelease())
                {
                    // The hook is dead, just get out of here
                    return Refs;
                }
            }
        }
        else
        {
            // Grab the owner reference before Release is called
            CBaseUnkHook* pOwner = Hook.m_pOwner;
            pOwner->_BaseAddRef();
            Refs = Hook.m_pOuter->Release();
            Hook.m_pQIHook->AfterRelease((long)Refs, (long)ppThis);
            if (0 == pOwner->_BaseRelease())
            {
                // The hook is dead, just get out of here
                return Refs;
            }
        }
    }
    else
    {
        Refs = Hook.m_pOuter->Release();
    }
    if (Refs)
    {
        // Only replace if object isn't dead
        *((void**)Hook.m_pOuter) = (void*)&Hook;
    }
    return Refs;
}
HRESULT _stdcall CBaseUnkHook::CHook::GetClassInfo(CHook** ppThis, ITypeInfo** ppTI)
{
    HRESULT hr;
    CHook& Hook = **ppThis;
    *((void**)Hook.m_pOuter) = Hook.m_lpVTableStart; 
    hr = Hook.m_pOuter->GetClassInfo(ppTI);
    *((void**)Hook.m_pOuter) = (void*)&Hook; 
    return hr;
}

ULONG _ElemCount(SAFEARRAY* pArray)
{
    // In general, these should be 1d arrays, but
    // multiple dimensions don't really hurt any.
    USHORT cDims = pArray->cDims;
    if (0 == cDims)
    {
        return 0;
    }

    ULONG cElems = 1;
    while (cDims--)
    {
        cElems *= pArray->rgsabound[cDims].cElements;
    }
    return cElems;
}
inline ULONG ElemCount(SAFEARRAY* pArray)
{
    return pArray ? _ElemCount(pArray) : 0;
}

// CUnkHookAggregate implementation
HRESULT CUnkHookAggregate::CreateInstance(IUnknown* pUnk, SAFEARRAY* pData, SAFEARRAY* pIIDs, UnknownHook** ppHook)
{
    *ppHook = NULL;
    ULONG cElems = ElemCount(pData);
    if (0 == cElems)
    {
        return E_INVALIDARG;
    }
    ULONG cIIDs = ElemCount(pIIDs);  // 0 OK here, IIDs not required.

    HRESULT hr;
    CUnkHookAggregate* pAggHook;
    BYTE* pAlloc;
    CAggregator::PreAllocData PAD((AggregateData*)pData->pvData, cElems, cIIDs, hr);

    if (FAILED(hr))
    {
        return hr;
    }

    if (NULL == (pAlloc = new BYTE[sizeof(CUnkHookAggregate) + PAD.GetExtraBytesCount()]))
    {
        return E_OUTOFMEMORY;
    }
    PAD.SetExtraBytesLocation(pAlloc + sizeof(CUnkHookAggregate));

    pAggHook = new ((void*)pAlloc) CUnkHookAggregate((AggregateData*)pData->pvData, cElems, cIIDs ? (IID*)pIIDs->pvData : NULL, PAD, hr, pUnk);
    if (SUCCEEDED(hr))
    {
        *ppHook = pAggHook;
        CUnkHookAggregate& Hook = **((CUnkHookAggregate**)ppHook);
        Hook.AddRef();
        Hook.m_Hook.HookUnk(reinterpret_cast<IProvideClassInfo*>(pUnk), (IQIARHook*)(IQIHook*)pAggHook, Hook.GetFlags(), true, false, &Hook);
    }
    else
    {
        delete pAggHook;
    }

    return hr;
}
STDMETHODIMP
CUnkHookAggregate::QIHook(VBGUID* iid, UnkHookFlags uhFlags, IUnknown** ppResult, IUnknown* Unknown)
{
    if (*ppResult || FastIsEqualIID(*(IID*)iid, IID_IUnknown))
    {
        return NOERROR;
    }

    IIDHookNode* pNodeToWrap = NULL;
    DoQIHook(*(IID*)iid, pNodeToWrap, (uhFlags & uhAfterQI) == uhAfterQI, ppResult);
    if (*ppResult && 
        pNodeToWrap)
    {
        IUnknown* pOldResult = *ppResult;

        // Use a NULL destructor because there is actually nothing to 
        // do when the aggregator is still alive.  Del.m_dwRefs is already 0,
        // and the _BlindDelegator clears its own members.
        // m_pfnDestroy is changed to PostMortemDestruct if the blind delegator
        // outlives the hook.
        _BlindDelegator::CreateDelegator(Unknown, pOldResult, NULL, &pNodeToWrap->BlindDel, NULL, ppResult);
        pOldResult->Release();
    }
    return NOERROR;
}
STDMETHODIMP
CUnkHookAggregate::MapIID(VBGUID *iid)
{
    return DoMapIID((IID*)iid);
}
STDMETHODIMP
CUnkHookAggregate::get_Flags(UnkHookFlags* puhFlags)
{
    *puhFlags = GetFlags();
    return NOERROR;
}
STDMETHODIMP
CUnkHookAggregate::put_Flags(UnkHookFlags uhFlags)
{
    // This is not allowed externally when an UnknownHook is
    // created by an AggregateUnknown call.
    return CTL_E_SETNOTSUPPORTED; // Facility control, error 383
}

CAggregator::PreAllocData::PreAllocData(AggregateData* pAggData, ULONG cAggElems, ULONG cIIDElems, HRESULT& hrLaunch)
{
    ULONG i;                  // Loop counter
    AggregateData* pCurData;  // Current aggdata for loop
    ULONG cElems;             // A local for counting elements

    ZeroMemory(this, sizeof(*this));

    // Initialize launch as successful
    hrLaunch = NOERROR;

    // Make a pass to get the required IID count
    i = cAggElems;
    pCurData = pAggData;
    while (i--)
    {
        ADFlags& Flags = pCurData->Flags;

        // If we have an object, then we always count the entry
        if (pCurData->pObject &&
            (0 == (Flags & adDontQuery)))
        {
            // Do some flag validation. You can't specify both of these.
            if ((Flags & (adWeakRefRaw | adWeakRefBalanced)) == (adWeakRefRaw | adWeakRefBalanced))
            {
                hrLaunch = E_INVALIDARG;
                return;        
            }

            if ((Flags & adPrimaryDispatch) &&
                !fHaveDispHook)
            {
                cIIDs++;
                fHaveDispHook = true;
            }
            if (Flags & adUseIIDs)
            {
                // Ignore LastIID if Last < First
                cElems = (pCurData->LastIID <= pCurData->FirstIID) ? 
                           1 : pCurData->LastIID - pCurData->FirstIID + 1;
                if (((Flags & adFullyResolved) && cElems != 1) ||
                    ((cElems + (ULONG)pCurData->FirstIID) > cIIDElems))
                {
                    hrLaunch = DISP_E_BADINDEX;
                    return;
                }
                cIIDs += cElems;
            }
            else
            {
                // Can't delay create without adUseIIDs
                if (Flags & adDelayCreation)
                {
                    hrLaunch = E_INVALIDARG;
                    return;
                }

                cUnfiltered++;
                if (Flags & adBeforeHookedUnknown)
                {
                    cUnfilteredBefore++;
                }
                if (Flags & adNoDelegator)
                {
                    cUnfilteredNoDelegator++;
                }
            }
        }
        // If we don't have an object, then this is an attempt to
        // change the primary dispatch to something other than the primary object.
        else if (pCurData->pObject == NULL &&
                 ((Flags & (adPrimaryDispatch | adUseIIDs)) == (adPrimaryDispatch | adUseIIDs)) &&
                 !fHaveDispHook)
        {
            // We only use the first IID
            if ((ULONG)pCurData->FirstIID >= cIIDElems)
            {
                hrLaunch = DISP_E_BADINDEX;
                return;
            }
            cIIDs++;
            fHaveDispHook = true;
        }
        else if (Flags & adMapIID)
        {
            if ((ULONG)pCurData->FirstIID >= cIIDElems ||
                (ULONG)pCurData->LastIID >= cIIDElems)
            {
                hrLaunch = DISP_E_BADINDEX;
                return;
            }
            cIIDMaps++;
            cIIDMapAndBlockElems++;
        }
        else if (Flags & adBlockIIDs)
        {
            // Ignore LastIID if Last < First
            cElems = (pCurData->LastIID <= pCurData->FirstIID) ? 
                       1 : pCurData->LastIID - pCurData->FirstIID + 1;
            if ((cElems + (ULONG)pCurData->FirstIID) > cIIDElems)
            {
                hrLaunch = DISP_E_BADINDEX;
                return;
            }
            cIIDBlocks += cElems;
            cIIDMapAndBlockElems++;
        }
        pCurData++;
    }

    // Figure out bytes count for the Unknown elements, as well as the
    // IID maps.  These are all of the items in CAggregator that are
    // created with the object instead of the FSMM.
    cExtraBytes = (cAggElems - cIIDMapAndBlockElems) * sizeof(CAggregator::CHoldUnk)
                   + cIIDMaps * sizeof(CAggregator::IIDMapEntry)
                   + cIIDBlocks * sizeof(IID);
}
CAggregator::CAggregator(AggregateData* pAggData, ULONG cAggElems, IID* pIIDs, PreAllocData& rPAD, HRESULT& hrLaunch, IUnknown* punkControlling)
      // Put anything here that would cause a problem
      // in the destructor if it weren't initialized.
    : m_pUnks(NULL),
      m_pIIDMaps(NULL),
      m_pBlockIIDs(NULL),
      m_pFSMMNodes(NULL)
{
    ULONG i;              // Loop counter
    AggregateData* pCurData;  // Current aggdata for loop
    ULONG cElems;             // A local for counting elements
    ULONG cFiltered;      // The number of encountered filtered objects.  Order in these doesn't matter, so they're tossed at the end.
    ULONG iCurEntry;      // The location in the m_pUnks array where the current element is stored
    IIDHookNode* pCurNode; // The node currently being worked on.
    IID* pCurIID;         // A loop iterator for items in the IID array
    IIDMapEntry* pCurMap; // The current position in the m_pIIDMaps array
    IID* pCurBlockIID;    // The current position in the m_pBlockIIDs array

    // Initialize launch as successful
    hrLaunch = NOERROR;

    m_cIIDMaps = rPAD.cIIDMaps;
    m_cBlockIIDs = rPAD.cIIDBlocks;
    m_cUnks = cAggElems - rPAD.cIIDMapAndBlockElems;

    // Start allocating our memory
    cElems = rPAD.cIIDs;
    if (i = rPAD.cUnfiltered - rPAD.cUnfilteredNoDelegator)
    {
        cElems += 2 * i;
        if (cElems < 4) cElems = 4;
    }
    if (cElems)
    {
        if (FAILED(hrLaunch = CFixedMemMgr::CreateInstance(sizeof(IIDHookNode), 
            // Try to alloc enough to get us through this routine up front,
            // plus two slots for each unfiltered item.  If this succeeds, then
            // we don't have to worry about FSMM->Alloc calls failing in this routine.
            cElems,
            &m_pFSMMNodes)))
        {
            return;
        }
    }

    if (rPAD.cExtraBytes)
    {
        // Zeroing is sufficient.  CHoldUnk and IIDMapNode
        // both need to be zero initialized.
        ZeroMemory(rPAD.pExtraBytes, rPAD.cExtraBytes);
    }

    if (m_cUnks)
    {
        //Change this block of code if CHoldUnk's constructor changes
        //CHoldUnk* pHoldUnk;   // A loop iterator for CHoldUnk items
        m_pUnks = (CHoldUnk*)rPAD.pExtraBytes;
        //i = m_cUnks;
        //pHoldUnk = m_pUnks;
        //while (i--)
        //{
        //    new ((void*)pHoldUnk) CHoldUnk;            
        //    pHoldUnk++;
        //}
        rPAD.pExtraBytes = (BYTE*)rPAD.pExtraBytes + m_cUnks * sizeof(CHoldUnk);
    }

    if (m_cIIDMaps)
    {
        // These are sufficiently constructed with the ZeroMemory on
        // the extra bytes
        pCurMap = m_pIIDMaps = (IIDMapEntry*)rPAD.pExtraBytes;
        rPAD.pExtraBytes = (BYTE*)rPAD.pExtraBytes + m_cIIDMaps * sizeof(IIDMapEntry);
    }

    if (m_cBlockIIDs)
    {
        pCurBlockIID = m_pBlockIIDs = (IID*)rPAD.pExtraBytes;
    }
    m_cUnfilteredBefore = m_cUnfilteredAfter = 0;

    // Now, actually fill the arrays
    i = cAggElems;
    pCurData = pAggData;
    cFiltered = 0;
    while (i--)
    {
        ADFlags& Flags = pCurData->Flags;
        if (pCurData->pObject)
        {
            if (Flags & adDontQuery)
            {
                // Just toss it near the end of the list.  This entry won't be referenced by
                // any IIDData, so it doesn't hurt anything.
                m_pUnks[m_cUnks - cFiltered - 1] = pCurData->pObject;
                cFiltered++;
            }
            else
            {
                if (Flags & adUseIIDs)
                {
                    // We have a filtered item (it gets sent only predefined IIDs)

                    // Find out how many IIDs we have.
                    // Ignore LastIID if Last < First
                    cElems = (pCurData->LastIID <= pCurData->FirstIID) ? 
                               1 : pCurData->LastIID - pCurData->FirstIID + 1;
                    
                    // We just throw this anywhere at the end of the array.
                    iCurEntry = m_cUnks - cFiltered - 1;
                    cFiltered++;
                    pCurIID = &pIIDs[pCurData->FirstIID];
                    IIDHookNodeList& rNodeList = (Flags & adBeforeHookedUnknown) ? m_BeforeNodes : m_AfterNodes;
                    while (cElems--)
                    {
                        // FSMM allocates the first chunk with the object,
                        // so this can't fail.
                        m_pFSMMNodes->Alloc((long*)&pCurNode);
                        pCurNode->iid = *pCurIID;
                        pCurNode->iUnk = (short)iCurEntry;
                        rNodeList.AddTail(pCurNode);
                        pCurIID++;
                    }

                    // The delayed creation flag is ignored if no IID's are specified, so
                    // we must check it in here.
                    if (Flags & adDelayCreation)
                    {
                        // If we're delaying creation, then there's some
                        // more work to do.
                        IDelayCreation* pDelay;
                        IUnknown* pTest = NULL;
                        HRESULT hr;
                        if (SUCCEEDED(hr = pCurData->pObject->QueryInterface(IID_IDelayCreation, (void**)&pDelay)))
                        {
                            m_pUnks[iCurEntry] = pDelay;
                            if (punkControlling &&
                                SUCCEEDED(hr = pDelay->QueryInterface(IID_IUnknown, (void**)&pTest)))
                            {
                                if (pTest == punkControlling)
                                {
                                    m_pUnks[iCurEntry].SetWeakRef();
                                    // Keep the reference on pDelay and release
                                    // a reference on the controlling IUnknown to
                                    // stay balanced.
                                    pTest->Release();
                                }
                                pTest->Release();
                            }
                            m_pUnks[iCurEntry].SetDelayedCreation();
                            if ((Flags & adDelayDontCacheResolved) == adDelayDontCacheResolved) // Two bit flag, need ==
                            {
                                m_pUnks[iCurEntry].SetKeepDelayed();
                            }
                            pDelay->Release();

                            if (Flags & adWeakRefBalanced)
                            {
                                m_pUnks[iCurEntry].SetExplicitWeakRef();
                            }
                            else if (Flags & adWeakRefRaw)
                            {
                                m_pUnks[iCurEntry].SetRawWeakRef();
                            }
                        }
                        if (FAILED(hr))
                        {
                            hrLaunch = DISP_E_TYPEMISMATCH;
                            return;
                        }
                    }

                    // Flag if we're fully resolved. This applies to items that
                    // are delay created as well as items with a single IID.
                    if (Flags & adFullyResolved)
                    {
                        m_pUnks[iCurEntry].SetFullyResolved();
                    }
                }
                else
                {
                    // This is an unfiltered item (it gets blindly sent IIDs)
                    if (Flags & adBeforeHookedUnknown)
                    {
                        iCurEntry = m_cUnfilteredBefore;
                        m_cUnfilteredBefore++;
                    }
                    else
                    {
                        iCurEntry = rPAD.cUnfilteredBefore + m_cUnfilteredAfter;
                        m_cUnfilteredAfter++;
                    }
                }

                // Share object assignment code for both filtered and blind cases.
                if (0 == (Flags & adDelayCreation))
                {
                    CHoldUnk& rHoldUnk = m_pUnks[iCurEntry];
                    if (Flags & adWeakRefRaw)
                    {
                        rHoldUnk.AssignRawWeakRef(pCurData->pObject);
                    }
                    else
                    {
                        rHoldUnk = pCurData->pObject;
                        if (Flags & adWeakRefBalanced)
                        {
                            IUnknown* punkControl = NULL;
                            pCurData->pObject->QueryInterface(IID_IUnknown, (void**)&punkControl);
                            if (punkControl)
                            {
                                punkControl->Release();
                                punkControl->Release();
                                rHoldUnk.SetWeakRef();
                            }
                        }
                    }
                }

                // See if the blind delegation has been
                // explicitly turned off.
                if (Flags & adNoDelegator)
                {
                    m_pUnks[iCurEntry].SetNoDelegator();
                }

                
                // If we found a primary dispatch request
                // during the initial pass and we haven't found
                // one yet, then deal with it the same as a 
                // before request.  An IDispatch can come on
                // a filtered or unfiltered object, and always
                // comes before.
                if (rPAD.fHaveDispHook &&
                    (Flags & adPrimaryDispatch))
                {
                    m_pFSMMNodes->Alloc((long*)&pCurNode);
                    pCurNode->iid = IID_IDispatch;
                    pCurNode->iUnk = (short)iCurEntry;
                    m_BeforeNodes.AddTail(pCurNode);
                    rPAD.fHaveDispHook = false;
                }
            }
        }
        else if (rPAD.fHaveDispHook &&
                 ((Flags & (adPrimaryDispatch | adUseIIDs)) == (adPrimaryDispatch | adUseIIDs)))
        {
            // This is a special case.  We're remapping the
            // primary dispatch of the controlling IUnknown to a secondary
            // interface on the controlling IUnknown.  This requires a
            // QI at this point, a WeakRef on the holder, and a flag
            // indicating that a QI should not be done during the QueryInterface
            // call.
            IUnknown* pResolved;
            if (SUCCEEDED(punkControlling->QueryInterface(pIIDs[pCurData->FirstIID], (void**)&pResolved)))
            {
                iCurEntry = m_cUnks - cFiltered - 1;
                cFiltered++;
                m_pFSMMNodes->Alloc((long*)&pCurNode);
                pCurNode->iid = IID_IDispatch;
                pCurNode->iUnk = (short)iCurEntry;
                m_BeforeNodes.AddTail(pCurNode);
                m_pUnks[iCurEntry].SetWeakRef();
                m_pUnks[iCurEntry].SetFullyResolved();
                m_pUnks[iCurEntry] = pResolved; // AddRef
                punkControlling->Release(); // Establish weakref state
                pResolved->Release(); // Balance QI call
                rPAD.fHaveDispHook = false;
            }
            else
            {
                hrLaunch = DISP_E_TYPEMISMATCH;
                return;
            }
        }
        else if (Flags & adMapIID)
        {
            pCurMap->IIDFrom = pIIDs[pCurData->FirstIID];
            pCurMap->IIDTo = pIIDs[pCurData->LastIID];
            pCurMap++;
        }
        else if (Flags & adBlockIIDs)
        {
            cElems = (pCurData->LastIID <= pCurData->FirstIID) ? 
                       1 : pCurData->LastIID - pCurData->FirstIID + 1;
            pCurIID = &pIIDs[pCurData->FirstIID];
            while (cElems--)
            {
                *pCurBlockIID++ = *pCurIID++;
            }
        }
        pCurData++;
    }
}
UnkHookFlags CAggregator::GetFlags()
{
    int Flags = 0;
    if (m_cUnfilteredBefore || m_BeforeNodes.pHead)
    {
        Flags |= uhBeforeQI;
    }
    if (m_cUnfilteredAfter || m_AfterNodes.pHead)
    {
        Flags |= uhAfterQI;
    }
    if (m_pIIDMaps || m_pBlockIIDs)
    {
        Flags |= uhMapIIDs;
    }

    return (UnkHookFlags)Flags;
}
CAggregator::~CAggregator()
{
    if (m_pUnks)
    {
        // Run the destructor by hand, these are
        // allocated with the item, so we can't
        // use delete[].
        CHoldUnk* pCur = m_pUnks;
        ULONG i = m_cUnks;
        while (i--)
        {
            pCur->~CHoldUnk();
            pCur++;
        }
    }
    if (m_pFSMMNodes)
    {
        OnTerminateAggregator(m_BeforeNodes);
        OnTerminateAggregator(m_AfterNodes);
        m_pFSMMNodes->Release();
    }
}
void CAggregator::OnTerminateAggregator(IIDHookNodeList& NodeList)
{
    IIDHookNode* pCur = NodeList.pHead;
    FixedSizeMemoryManager* pFSMM = m_pFSMMNodes;
    while (pCur)
    {
        if (pCur->BlindDel.m_dwRefs)
        {
            (pCur->pFSMMOwner = pFSMM)->AddRef();
            pCur->BlindDel.m_pfnDestroy = CAggregator::IIDHookNode::PostMortemDestruct;
        }
        pCur = pCur->pNext;
    }
}
ULONG _stdcall CAggregator::IIDHookNode::PostMortemDestruct(_BlindDelegator* pBlindDel, IUnknown** ppInner, IUnknown** ppOuter)
{
    IIDHookNode& Node = *(IIDHookNode*)pBlindDel;
    Node.pFSMMOwner->Release();
    return 0; // Ignored if ppOuter not freed
}

HRESULT CAggregator::DoQIHook(REFIID riid, CAggregator::IIDHookNode*& pNodeToWrap, bool fAfter, IUnknown** ppResult)
{
    // First, walk the filtered before values
    IIDHookNodeList& NodesList = fAfter ? m_AfterNodes : m_BeforeNodes;
    IIDHookNode* pCurNode = NodesList.pHead;
    HRESULT hrTest;
    short i, iEnd;

    while (pCurNode)
    {
        IIDHookNode& rNode = *pCurNode;
        if (FastIsEqualIID(riid, rNode.iid))
        {
            if (rNode.BlindDel.m_dwRefs)
            {
                (*ppResult = (IUnknown*)(&rNode.BlindDel.m_pVTable))->AddRef();
                // Don't set pNodeToWrap, there's nothing more to be done
                // with this value.
                return NOERROR;
            }
            else
            {
                CHoldUnk& rUnkHold = m_pUnks[rNode.iUnk];
                IUnknown* pResolved = NULL;
                if (!rUnkHold.GetDelayedCreation() || SUCCEEDED(rUnkHold.ResolveDelayedCreation(riid, &pResolved)))
                {
                    if (pResolved)
                    {
                        // In this case, rUnkHold still holds the IDelayedCreation
                        // interface reference, so we QI the returned object.
                        if (rUnkHold.GetFullyResolved())
                        {
                            *ppResult = pResolved;
                            hrTest = NOERROR;
                        }
                        else
                        {
                            hrTest = pResolved->QueryInterface(riid, (void**)ppResult);
                            pResolved->Release();
                        }
                    }
                    else if (rUnkHold.GetFullyResolved())
                    {
                        // No QI is necessary (generally an IDispatch call)
                        (*ppResult = rUnkHold.operator->())->AddRef();
                        hrTest = NOERROR;
                    }
                    else
                    {
                        // Do a normal QueryInterface
                        hrTest = rUnkHold->QueryInterface(riid, (void**)ppResult);
                    }

                    if (SUCCEEDED(hrTest))
                    {
                        if (!rUnkHold.GetNoDelegator())
                        {
                            pNodeToWrap = pCurNode;
                        }
                        return NOERROR;
                    }
                }
            }
        }
        pCurNode = pCurNode->pNext;
    }

    // Now try all the unfiltered objects
    // Before loop: for (i = 0; i < m_cUnfilteredBefore; i++)
    // After loop:  for (i = m_cUnfilteredBefore; i < m_cUnfilteredBefore + m_cUnfilteredAfter; i++)
    for (fAfter ? (i = m_cUnfilteredBefore, iEnd = m_cUnfilteredBefore + m_cUnfilteredAfter) :
                  (i = 0, iEnd = m_cUnfilteredBefore);
         i < iEnd;
         i++)
    {
        if (SUCCEEDED(m_pUnks[i]->QueryInterface(riid, (void**)ppResult)))
        {
            if (!m_pUnks[i].GetNoDelegator())
            {
                IIDHookNode* pNewNode;
                if (FAILED(hrTest = m_pFSMMNodes->Alloc((long*)&pNewNode)))
                {
                    (*ppResult)->Release();
                    *ppResult = NULL;
                    return hrTest; // We're out of memory, no point continuing
                }
                // FSMM does a ZeroMemory, so we don't have
                // much to initialize.
                NodesList.AddTail(pNewNode);
                pNewNode->iid = riid;
                pNewNode->iUnk = i;
                pNodeToWrap = pNewNode;
            }
            return NOERROR;
        }
    }
    return E_NOINTERFACE;
}
HRESULT CAggregator::DoMapIID(IID* pIID)
{
    ULONG i;

    // First check the block list
    if (i = m_cBlockIIDs)
    {
        IID* pBlockIID = m_pBlockIIDs;
        while (i--)
        {
            if (FastIsEqualIID(*pIID, *pBlockIID))
            {
                return E_NOINTERFACE;
            }
            pBlockIID++;
        }
    }


    if (i = m_cIIDMaps)
    {
        IIDMapEntry* pCur = m_pIIDMaps;
        while (i--)
        {
            if (FastIsEqualIID(*pIID, pCur->IIDFrom))
            {
                *pIID = pCur->IIDTo;
                break;
            }
            pCur++;
        }
    }
    return NOERROR;
}
HRESULT CAggregator::CHoldUnk::ResolveDelayedCreation(REFIID riid, IUnknown** ppResolved)
{
    //ASSERT(m_fDelayedCreation);
    //ASSERT(ppResolved && *ppResolved == NULL);
    HRESULT hr;
    IUnknown* pUnkResolved = NULL;
    if (SUCCEEDED(hr = ((IDelayCreation*)m_pUnk)->Create((VBGUID*)&riid, &pUnkResolved)))
    {
        if (pUnkResolved)
        {
            if (m_fKeepDelayed)
            {
                *ppResolved = pUnkResolved;
            }
            else
            {
                if (m_fWeakRef)
                {
                    // Balance the missing ref on the controlling IUnknown
                    IUnknown* punkControl;
                    m_pUnk->QueryInterface(IID_IUnknown, (void**)&punkControl);
                }
                m_pUnk->Release();
                m_fWeakRef = m_fDelayedCreation = false;
                m_pUnk = pUnkResolved; //Already AddRef'ed by Create, don't use operator=

                // This should be held as a weak reference. (Re)establish the m_fWeakRef
                // flag.
                if (m_fWeakRefExplicit)
                {
                    IUnknown* punkControl = NULL;
                    m_pUnk->QueryInterface(IID_IUnknown, (void**)&punkControl);
                    if (punkControl)
                    {
                        // Balance the QI call
                        punkControl->Release();
                        // Enter the WeakRef state
                        punkControl->Release();
                        m_fWeakRef = true;
                    }
                }
            }
        }
        else
        {
            hr = E_NOINTERFACE;
        }
    }
    return hr;
}
// CAggregateObject implementation
HRESULT CAggregateObject::CreateInstance(SAFEARRAY* pAggData, SAFEARRAY* pIIDs, IUnknown** ppOwner, IUnknown** ppUnk)
{
    *ppUnk = NULL;
    ULONG cElems = ElemCount(pAggData);
    if (0 == cElems)
    {
        return E_INVALIDARG;
    }
    ULONG cIIDs = ElemCount(pIIDs);  // 0 OK here, IIDs not required.

    HRESULT hr;
    CAggregateObject* pAggObj;
    BYTE* pAlloc;
    CAggregator::PreAllocData PAD((AggregateData*)pAggData->pvData, cElems, cIIDs, hr);

    if (FAILED(hr))
    {
        return hr;
    }

    if (NULL == (pAlloc = new BYTE[sizeof(CAggregateObject) + PAD.GetExtraBytesCount()]))
    {
        return E_OUTOFMEMORY;
    }
    PAD.SetExtraBytesLocation(pAlloc + sizeof(CAggregateObject));

    pAggObj = new ((void*)pAlloc) CAggregateObject((AggregateData*)pAggData->pvData, cElems, cIIDs ? (IID*)pIIDs->pvData : NULL, PAD, hr);

    if (SUCCEEDED(hr))
    {
        *ppUnk = pAggObj;
        pAggObj->m_UHFlags = pAggObj->GetFlags();
        if (ppOwner)
        {
            *ppOwner = pAggObj;
            pAggObj->m_ppThis = ppOwner;
        }
        pAggObj->AddRef();
    }
    else
    {
        REMOVE_OBJECT; // Balances ADD_OBJECT in constructor, there's no dtor
        delete pAggObj;
    }
    return hr;
}
STDMETHODIMP
CAggregateObject::QueryInterface(REFIID riid, void** ppv)
{
    if (FastIsEqualIID(riid, IID_IUnknown))
    {
        m_dwRefs++;
        *ppv = this;
        return NOERROR;
    }
    else
    {
        IUnknown* pResult = NULL;
        IIDHookNode* pNodeToWrap = NULL;
        IID* piidUse = const_cast<IID*>(&riid);
        IID IIDMapped;
        if (m_UHFlags & uhMapIIDs)
        {
            IIDMapped = riid;
            if (FAILED(DoMapIID(&IIDMapped)))
            {
                goto MapIIDFailure;
            }
            piidUse = &IIDMapped;
        }
        if (m_UHFlags & uhBeforeQI)
        {
            DoQIHook(*piidUse, pNodeToWrap, false, &pResult);
        }
        if ((NULL == pResult) && (m_UHFlags & uhAfterQI))
        {
            DoQIHook(*piidUse, pNodeToWrap, true, &pResult);
        }
        if (pResult && 
            pNodeToWrap)
        {
            IUnknown* pOldResult = pResult;
            _BlindDelegator::CreateDelegator(this, pOldResult, NULL, &pNodeToWrap->BlindDel, NULL, &pResult);
            pOldResult->Release();
        }
MapIIDFailure:
        return (*ppv = pResult) ? NOERROR : E_NOINTERFACE;
    }
}
